import sqlite3
import xml.etree.ElementTree as ET
from sqlite3 import Connection, Cursor
from dataclasses import dataclass


class Node:
    def __init__(self, node_id: str, *, value=None, parent=None, source=None, target=None) -> None:
        self.node_id = node_id
        self.value = value
        self.parent = parent
        self.source = source
        self.target = target


@dataclass
class Diagram:
    diagram_id: str
    name: str
    # 父状态的ID
    parent_state_id: str = None

    def is_subdiagram(self):
        return bool(self.parent_state_id)


@dataclass
class DiagramElement:
    model_element_id: str
    shape_type: str

    def is_state(self):
        return self.shape_type == 'State2'

    def is_transition(self):
        return self.shape_type == 'Transition2'

    def is_initial_pseudo_state(self):
        return self.shape_type == 'InitialPseudoState'

@dataclass
class ModelElement:
    model_id: str
    name: str
    definition: str

    def get_from(self):
        from_prefix = 'fromModel=<'
        from_suffix = '>;'
        
        from_index_start = self.definition.find(from_prefix) + len(from_prefix)
        from_index_end = self.definition.find(from_suffix, from_index_start)

        ret = self.definition[from_index_start:from_index_end]
        if ':' in ret:
            return ret.split(':')[1]
        else:
            return ret

    def get_to(self):
        to_prefix = 'toModel=<'
        to_suffix = '>;'

        to_index_start = self.definition.find(to_prefix) + len(to_prefix)
        to_index_end = self.definition.find(to_suffix, to_index_start)
        
        ret = self.definition[to_index_start:to_index_end]
        if ':' in ret:
            return ret.split(':')[1]
        else:
            return ret

    def get_transitions(self):
        from_id = self.get_from()
        to_id = self.get_to()

        if ':' in from_id:
            from_id = from_id.split(':')[1]

        if ':' in to_id:
            to_id = to_id.split(':')[1]

        return (from_id, to_id)


def get_diagrams(conn: Connection):
    c: Cursor = conn.cursor()
    ret = []
    for r in c.execute('SELECT ID, NAME, PARENT_MODEL_ID FROM DIAGRAM'):
        ret.append(Diagram(r[0], r[1], r[2]))

    c.close()

    return ret


def get_diagram_models(conn: Connection, diagram_id: str):
    c: Cursor = conn.cursor()
    ret = []

    for r in c.execute("SELECT MODEL_ELEMENT_ID, SHAPE_TYPE FROM DIAGRAM_ELEMENT WHERE DIAGRAM_ID=?", (diagram_id, )):
        ret.append(DiagramElement(r[0], r[1]))

    c.close()

    return ret


def get_model_element(conn: Connection, model_id: str) -> ModelElement:
    c: Cursor = conn.cursor()

    r = c.execute('SELECT ID, NAME, DEFINITION FROM MODEL_ELEMENT WHERE ID=?', (model_id, ))

    r = r.fetchone()

    c.close()

    return ModelElement(
        r[0], 
        r[1], 
        r[2].decode('utf8'))


@dataclass
class StateMachine:
    name: str
    init_state: str
    states: list
    transitions: list

@dataclass
class State:
    name: str
    sub_machine: StateMachine = None


@dataclass
class Transition:
    from_state: str
    to_state: str
    trigger: str

    @staticmethod
    def parse_transition(from_state: str, to_state: str, name: str) -> list:
        if name:
            return [Transition(from_state, to_state, i) for i in name.split('/') if i]
        else:
            return []

    def transition_method_name(self, camel: bool = True):
        if not camel:
            return f'from_{self.from_state}_to_{self.to_state}_by_{self.trigger}'
        else:
            return f'from{self.from_state.title()}To{self.to_state.title()}By{self.trigger.title()}'

    def __repr__(self):
        return f'{self.from_state} -> ({self.transition_method_name()})-> {self.to_state}, trigger: {self.trigger}'


def get_machine(conn: Connection):
    diagrams = get_diagrams(conn)

    # 状态名 -> machine
    sub_machine_dict = {}

    machines = []

    state_dict = {}

    for d in diagrams:
        models = get_diagram_models(conn, d.diagram_id)
        db_transitions = [get_model_element(conn, t.model_element_id) for t in models if t.is_transition()]
        pseudo_initial_state = [t for t in models if t.is_initial_pseudo_state()]
        db_states = [t for t in models if t.is_state()]

        # 找到第一个状态
        pseudo_state_id = pseudo_initial_state[0].model_element_id

        transitions = []

        for t in db_transitions:
            from_id = t.get_from()
            to_id = t.get_to()
            if from_id == pseudo_state_id:
                init_state_id = t.get_to()
            else:
                transitions += Transition.parse_transition(
                    get_model_element(conn, from_id).name,
                    get_model_element(conn, to_id).name,
                    t.name)

        states = []
        for s in db_states:
            state_in_db = get_model_element(conn, s.model_element_id)
            state = State(state_in_db.name)
            states.append(state)
            state_dict[state.name] = state
            if init_state_id == s.model_element_id:
                init_state_name = state_in_db.name

        machine = StateMachine(d.name, init_state_name, states, transitions)

        if d.is_subdiagram():
            s = get_model_element(conn, d.parent_state_id)
            sub_machine_dict[s.name] = machine

        machines.append(machine)

        for m in machines:
            for s in m.states:
                if s.name in sub_machine_dict.keys():
                    s.sub_machine = sub_machine_dict[s.name]

    for m in machines:
        if not (m.name in [m.name for m in sub_machine_dict.values()]):
            return (state_dict, m)


def dump_transition_methods(machine: StateMachine) -> list:
    ret = []

    for t in machine.transitions:
        ret.append(t.transition_method_name())

    for s in machine.states:
        if not (s.sub_machine is None):
            ret += dump_transition_methods(s.sub_machine)

    return ret


def dump_transition_triggers(machine: StateMachine) -> list:
    ret = []

    for t in machine.transitions:
        ret.append(t.trigger)

    for s in machine.states:
        if not (s.sub_machine is None):
            ret += dump_transition_triggers(s.sub_machine)

    return set(ret)


def to_java_repr(s_dict: dict, target_machine: StateMachine, class_name, package_name) -> str:
    import misc_java
    ret = ""

    ret += misc_java.header(package_name, class_name)

    ret += misc_java.states_header()
    ret += misc_java.states_repr(target_machine)
    ret += misc_java.states_footer()

    ret += misc_java.transitions_header()
    ret += misc_java.transitions_repr(target_machine)
    ret += misc_java.transitions_footer()

    ret += misc_java.context_annotation()

    ret += misc_java.class_declaration_start(class_name)
    ret += misc_java.states_declaration(s_dict.keys())
    ret += misc_java.events_declaration(dump_transition_triggers(target_machine))
    ret += '\r\n'
    ret += misc_java.transition_methods_repr(dump_transition_methods(target_machine))
    ret += misc_java.class_declaration_end()

    return ret


def to_kotlin_repr(s_dict: dict, target_machine: StateMachine, class_name, package_name) -> str:
    import misc_kotlin
    ret = ''

    ret += misc_kotlin.header(package_name, class_name)

    ret += misc_kotlin.states_header()
    ret += misc_kotlin.states_repr(target_machine, class_name)
    ret += misc_kotlin.states_footer()

    ret += misc_kotlin.transitions_header()
    ret += misc_kotlin.transitions_repr(target_machine, class_name)
    ret += misc_kotlin.transitions_footer()

    ret += misc_kotlin.context_annotation()

    ret += misc_kotlin.class_declaration_start(class_name)
    ret += misc_kotlin.transition_methods_repr(dump_transition_methods(target_machine))
    ret += '\r\n'
    ret += misc_kotlin.companion_declaration_start()
    ret += misc_kotlin.states_declaration(s_dict.keys())
    ret += misc_kotlin.events_declaration(dump_transition_triggers(target_machine))
    ret += misc_kotlin.companion_declaration_end()

    ret += misc_kotlin.class_declaration_end()

    return ret

def to_rust_repr(s_dict: dict, target_machine: StateMachine, class_name, package_name) -> str:
    import misc_rust

    ret = ''

    ret += misc_rust.header(package_name, class_name)

    ret += misc_rust.states_header()
    ret += misc_rust.states_repr(target_machine, class_name)
    ret += misc_rust.states_footer()

    ret += misc_rust.transitions_header()
    ret += misc_rust.transitions_repr(target_machine, class_name)
    ret += misc_rust.transitions_footer()

    ret += misc_rust.context_annotation()

    ret += misc_rust.class_declaration_start(class_name)
    ret += misc_rust.transition_methods_repr(dump_transition_methods(target_machine))
    ret += '\r\n'
    ret += misc_rust.companion_declaration_start()
    ret += misc_rust.states_declaration(s_dict.keys())
    ret += misc_rust.events_declaration(dump_transition_triggers(target_machine))
    ret += misc_rust.companion_declaration_end()

    ret += misc_rust.class_declaration_end()

    return ret



def to_python_repr(s_dict: dict, target_machine: StateMachine, class_name, package_name) -> str:
    import misc_python

    ret = misc_python.header()

    ret += misc_python.states_declaration(s_dict.keys())
    ret += misc_python.event_declaration(dump_transition_triggers(target_machine))

    ret += misc_python.class_declaration(class_name, s_dict.keys())

    ret += misc_python.transitions_repr(target_machine, class_name)

    ret += misc_python.methods_repr(dump_transition_methods(target_machine))

    return ret


def fsm_repr(s_dict: dict, target_machine: StateMachine, class_name, package_name, lang='java'):
    if lang == 'java':
        return to_java_repr(s_dict, target_machine, class_name, package_name)
    elif lang == 'kotlin':
        return to_kotlin_repr(s_dict, target_machine, class_name, package_name)
    elif lang == 'python':
        return to_python_repr(s_dict, target_machine, class_name, package_name)
    elif lang == 'rust':
        return to_rust_repr(s_dict, target_machine, class_name, package_name)


def get_drawio_machine(file: str):
    element = ET.parse(file).find('diagram').find('mxGraphModel').find('root')
    transition_tuples = [] #(id, source, target)
    node_dict = {}
    for e in element.iter():
        if e.tag != 'mxCell' or 'id' not in e.attrib.keys() or e.attrib['id'] == '0' or e.attrib['id'] == '1':
            continue

        node_dict[e.attrib['id']] = Node(
            e.attrib['id'], 
            value=e.attrib.get('value'), 
            parent=e.attrib.get('parent'), 
            source=e.attrib.get('source'), 
            target=e.attrib.get('target'))

        if 'source' in e.attrib.keys() and 'target' in e.attrib.keys():
            transition_tuples.append((e.attrib['id'], e.attrib['source'], e.attrib['target']))


    transition_node_tuples = [] # from, to, event

    
    state_set = set()

    start_states = set()
    end_states = set()

    for t in transition_tuples:
        print(f'transition: {t}')
        node_id, source_id, target_id = t
        for n in node_dict.values():
            if n.parent == node_id:
                event_node = n
                break


        source_node = node_dict[source_id]
        target_node = node_dict[target_id]

        start_states.add(source_node.value)
        end_states.add(target_node.value)
        
        transition_node_tuples.append((source_node, target_node, event_node))
        state_set.add(event_node.value)

    s_dict = {}
    for s in state_set:
        s_dict[s] = State(s)

    
    for s in start_states:
        if s not in end_states:
            start_state = s
            break

    transition_list = []
    for t in transition_node_tuples:
        source_node, target_node, event_node = t
        trans = Transition(source_node.value, target_node.value, event_node.value)
        transition_list.append(trans)

    machine = StateMachine(file.split('.')[0], start_state, list(map(lambda s: State(s), state_set)), transition_list)
    return (s_dict, machine)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='state machine code generator')
    parser.add_argument('-f', '--file', type=str)
    parser.add_argument('-c', '--classname', type=str, default='HelloMachine')
    parser.add_argument('-p', '--packagename', type=str, default='com.example.www')
    parser.add_argument('-l', '--language', type=str, default='kotlin')

    args = parser.parse_args()

    if args.file.endswith('xml'):
        state_dict, machine = get_drawio_machine(args.file)
    else:
        state_dict, machine = get_machine(sqlite3.connect(args.file))


    print(fsm_repr(state_dict, machine, args.classname, args.packagename, args.language))

